package com.practice.tree;

public class RedBlack {
	public int numTwists(int[] keys) {
		int count = 0;
		RBTree tree = new RBTree();
		for (int key : keys) {
			count += tree.insert(key);
		}
		return count;
	}
	
	private class RBTree {
		private Node		root;
		
		/**
		 * Returns number of twists needed.
		 * 
		 * @param key
		 * @return
		 */
		public int insert(int key) {
			Node node = newNode(key);
			return (node.parent!=null && node.parent.color==Node.RED)? twist(node) : 0;
		}
		
		private int twist(Node node) {
			Node parent = node.parent;
			Node grandParent = parent.parent;
			
			if (grandParent == null) {
				parent.color = Node.BLACK;
				return 0;
			}
			
			Node greatParent = grandParent.parent;
			
			Node y = null;
			if (parent.left == node && grandParent.left == parent) {
				grandParent.setLeft(parent.right);
				parent.setRight(grandParent);
				parent.parent = grandParent.parent;
				grandParent.parent = parent;
				
				node.color = Node.BLACK;
				
				y = parent;
			}
			else if (parent.right == node && grandParent.left == parent) {
				parent.setRight(node.left);
				grandParent.setLeft(node.right);
				node.setLeft(parent);
				node.setRight(grandParent);
				
				parent.color = Node.BLACK;
				
				y = node;
			}
			else if (parent.left == node && grandParent.right == parent) {
				grandParent.setRight(node.left);
				parent.setLeft(node.right);
				node.setLeft(grandParent);
				node.setRight(parent);
				
				parent.color = Node.BLACK;
				
				y = node;
			}
			else {	// parent.right == node && grandParent.right == parent
				grandParent.setRight(parent.left);
				parent.setLeft(grandParent);
				
				node.color = Node.BLACK;
				
				y = parent;
			}
			
			if (grandParent == root) {
				root = y;
				root.parent = null;
				return 1;
			}
			else {
				if (greatParent.left == grandParent) {
					greatParent.setLeft(y);
				}
				else {
					greatParent.setRight(y);
				}
			}
			
			return greatParent.color==Node.RED? twist(y)+1 : 1;
		}
		
		private Node newNode(int key) {
			Node node = new Node();
			node.key = key;
			node.color = Node.RED;
			Node parent = findParent(root, key); 
			if (parent == null) {
				root = node;
				root.color = Node.BLACK;
				return node;
			}
			
			if (parent.key > node.key) {
				parent.setLeft(node);
			}
			else {
				parent.setRight(node);
			}
			
			return node;
		}
		
		private Node findParent(Node parent, int key) {
			if (parent == null) return null;
			if (parent.key > key) {
				return parent.left==null? parent : findParent(parent.left, key);
			}
			else {
				return parent.right==null? parent : findParent(parent.right, key);
			}
		}
	}
	
	private class Node {
		public final static int		RED 	= 0;
		public final static int		BLACK 	= 1;
		
		private Node		left;
		private Node		right;
		private Node		parent;
		private int			key;
		private int			color;
		
		public void setLeft(Node left) {
			this.left = left;
			if (left != null) left.parent = this;
		}
		
		public void setRight(Node right) {
			this.right = right;
			if (right != null) right.parent = this;
		}
	}
}
